{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Object Detection using TAO DetectNet_v2\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\"> "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## What is DetectNet-V2?\n",
    "\n",
    "DetectNet_v2, also known as GridBox object detection, is a highly optimized CNN based object detection model, that does bounding-box regression over a uniform grid on the input image. The gridbox system divides an input image into a uniform grid that predicts four normalized bounding-box parameters (xc, yc, w, h) and confidence value per output class.\n",
    "\n",
    "The raw normalized bounding-box and confidence detection are thesholded and then post-processed by a clustering algorithm such as DBSCAN, NMS or a HYBRID (DBSCAN + NMS) to produce the final bounding-box coordinates and category labels.\n",
    "\n",
    "### Sample output predictions from a trained DetectNet_v2 model\n",
    "\n",
    "<img align=\"center\" src=\"https://miro.medium.com/v2/resize:fit:720/0*YaQDIKR4gRbP2-by\" width=\"960\">\n",
    "\n",
    "<img align=\"center\" src=\"https://miro.medium.com/v2/resize:fit:720/0*eY1qluSyldYl9qDw\" width=\"960\">"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained resnet18 model and train a ResNet-18 DetectNet_v2 model on the KITTI dataset\n",
    "* Prune the trained detectnet_v2 model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Export the pruned model\n",
    "* Quantize the pruned model using QAT\n",
    "* Run Inference on the trained model\n",
    "* Export the pruned, quantized and retrained model to a .onnx file for deployment to DeepStream\n",
    "* Run inference on the exported. etlt model to verify deployment using TensorRT\n",
    "\n",
    "At the end of this notebook, you will have a trained and optimized `detectnet_v2` model that you\n",
    "may deploy via [Triton](https://github.com/NVIDIA-AI-IOT/tao-toolkit-triton-apps) or [DeepStream](https://developer.nvidia.com/deepstream-sdk).\n",
    "\n",
    "### Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of Object Detection using DetectNet_v2 in the Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "    1. [Third Party MLOPS integration](#head-0-1)\n",
    "1. [Install the TAO Launcher](#head-1)\n",
    "1. [Prepare dataset and pre-trained model](#head-2)\n",
    "    1. [Download the dataset](#head-2-1)\n",
    "    1. [Verify downloaded dataset](#head-2-2)\n",
    "    1. [Prepare tfrecords from kitti format dataset](#head-2-3)\n",
    "    2. [Download pre-trained model](#head-2-4)\n",
    "2. [Provide training specification](#head-3)\n",
    "3. [Run TAO training](#head-4)\n",
    "4. [Evaluate trained models](#head-5)\n",
    "5. [Prune trained models](#head-6)\n",
    "6. [Retrain pruned models](#head-7)\n",
    "7. [Evaluate retrained model](#head-8)\n",
    "8. [Visualize inferences](#head-9)\n",
    "9. [Model Export](#head-10)\n",
    "10. [Verify Deployed Model](#head-11)\n",
    "    1. [Inference using TensorRT engine](#head-11-1)\n",
    "11. [QAT workflow](#head-12)\n",
    "    1. [Convert pruned model to QAT and retrain](#head-12-1)\n",
    "    2. [Evaluate QAT converted model](#head-12-2)\n",
    "    3. [Export QAT trained model to int8](#head-12-3)\n",
    "    4. [Evaluate a QAT trained model using the exported TensorRT engine](#head-12-4)\n",
    "    5. [Inference using QAT engine](#head-12-5)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "The following notebook requires the user to set an env variable called the `$LOCAL_PROJECT_DIR` as the path to the users workspace. Please note that the dataset to run this notebook is expected to reside in the `$LOCAL_PROJECT_DIR/data`, while the TAO experiment generated collaterals will be output to `$LOCAL_PROJECT_DIR/detectnet_v2`. More information on how to set up the dataset and the supported steps in the TAO workflow are provided in the subsequent cells.\n",
    "\n",
    "*Note: Please make sure to remove any stray artifacts/files from the `$USER_EXPERIMENT_DIR` or `$DATA_DOWNLOAD_DIR` paths as mentioned below, that may have been generated from previous experiments. Having checkpoint files etc may interfere with creating a training graph for a new experiment.*\n",
    "\n",
    "*Note: This notebook currently is by default set up to run training using 1 GPU. To use more GPU's please update the env variable `$NUM_GPUS` accordingly*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Setting up env variables for cleaner command line commands.\n",
    "import os\n",
    "\n",
    "%env NUM_GPUS=1\n",
    "%env USER_EXPERIMENT_DIR=/workspace/tao-experiments/detectnet_v2\n",
    "%env DATA_DOWNLOAD_DIR=/workspace/tao-experiments/data\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=~/tao-samples/detectnet_v2\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "# The dataset expected to be present in $LOCAL_PROJECT_DIR/data, while the results for the steps\n",
    "# in this notebook will be stored at $LOCAL_PROJECT_DIR/detectnet_v2\n",
    "# !PLEASE MAKE SURE TO UPDATE THIS PATH!.\n",
    "\n",
    "os.environ[\"LOCAL_PROJECT_DIR\"] = FIXME\n",
    "\n",
    "os.environ[\"LOCAL_DATA_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"data\"\n",
    ")\n",
    "os.environ[\"LOCAL_EXPERIMENT_DIR\"] = os.path.join(\n",
    "    os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()),\n",
    "    \"detectnet_v2\"\n",
    ")\n",
    "\n",
    "# Make the experiment directory \n",
    "! mkdir -p $LOCAL_EXPERIMENT_DIR\n",
    "\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"LOCAL_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "%env SPECS_DIR=/workspace/tao-experiments/detectnet_v2/specs\n",
    "CLEARML_LOGGED_IN = False\n",
    "WANDB_LOGGED_IN = False\n",
    "\n",
    "# Showing list of specification files.\n",
    "!ls -rlt $LOCAL_SPECS_DIR"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Third Party MLOPS integration <a class=\"anchor\" id=\"head-0-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment this code block to enable wandb integration with TAO.\n",
    "# # you can get the wandb api key by logging into https://wandb.ai\n",
    "# !pip install wandb\n",
    "\n",
    "# os.environ[\"WANDB_API_KEY\"] = FIXME\n",
    "# import wandb\n",
    "# WANDB_LOGGED_IN = wandb.login()\n",
    "# if WANDB_LOGGED_IN:\n",
    "#     print(\"WANDB successfully logged in.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment the code block in this cell to enable clearml integration with TAO.\n",
    "# !pip3 install clearml\n",
    "\n",
    "# # Copy the data from the clearml website when generating credentials through the clearml webapp.\n",
    "# # Please choose the jupyter notebook tab under the clearml credentials and copy the data as is\n",
    "# # or fill the values to the required env variables.\n",
    "\n",
    "# os.environ[\"CLEARML_WEB_HOST\"] = FIXME\n",
    "# os.environ[\"CLEARML_API_HOST\"] = FIXME\n",
    "# os.environ[\"CLEARML_FILES_HOST\"] = FIXME\n",
    "# os.environ[\"CLEARML_API_ACCESS_KEY\"] = FIXME\n",
    "# os.environ[\"CLEARML_API_SECRET_KEY\"] = FIXME\n",
    "\n",
    "# CLEARML_LOGGED_IN=True"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The cell below maps the project directory on your local host to a workspace directory in the TAO docker instance, so that the data and the results are mapped from in and out of the docker. For more information please refer to the [launcher instance](https://docs.nvidia.com/tao/tao-toolkit/text/tao_launcher.html) in the user guide.\n",
    "\n",
    "When running this cell on AWS, update the drive_map entry with the dictionary defined below, so that you don't have permission issues when writing data into folders created by the TAO docker.\n",
    "\n",
    "```json\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "            # Mapping the data directory\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "                \"destination\": \"/workspace/tao-experiments\"\n",
    "            },\n",
    "            # Mapping the specs directory.\n",
    "            {\n",
    "                \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "                \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "            },\n",
    "        ],\n",
    "    \"DockerOptions\": {\n",
    "        \"user\": \"{}:{}\".format(os.getuid(), os.getgid())\n",
    "    }\n",
    "}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "\n",
    "# Define the dictionary with the mapped drives\n",
    "drive_map = {\n",
    "    \"Mounts\": [\n",
    "        # Mapping the data directory\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "            \"destination\": \"/workspace/tao-experiments\"\n",
    "        },\n",
    "        # Mapping the specs directory.\n",
    "        {\n",
    "            \"source\": os.environ[\"LOCAL_SPECS_DIR\"],\n",
    "            \"destination\": os.environ[\"SPECS_DIR\"]\n",
    "        },\n",
    "    ],\n",
    "    \"DockerOptions\":{\n",
    "        \"user\": f\"{os.getuid()}:{os.getgid()}\"\n",
    "    }\n",
    "}\n",
    "\n",
    "if CLEARML_LOGGED_IN:\n",
    "    if \"Envs\" not in drive_map.keys():\n",
    "        drive_map[\"Envs\"] = []\n",
    "    drive_map[\"Envs\"].extend([\n",
    "        {\n",
    "            \"variable\": \"CLEARML_WEB_HOST\",\n",
    "            \"value\": os.getenv(\"CLEARML_WEB_HOST\")\n",
    "        },\n",
    "        {\n",
    "            \"variable\": \"CLEARML_API_HOST\",\n",
    "            \"value\": os.getenv(\"CLEARML_API_HOST\")\n",
    "        },\n",
    "        {\n",
    "            \"variable\": \"CLEARML_FILES_HOST\",\n",
    "            \"value\": os.getenv(\"CLEARML_FILES_HOST\")\n",
    "        },\n",
    "        {\n",
    "            \"variable\": \"CLEARML_API_ACCESS_KEY\",\n",
    "            \"value\": os.getenv(\"CLEARML_API_ACCESS_KEY\")\n",
    "        },\n",
    "        {\n",
    "            \"variable\": \"CLEARML_API_SECRET_KEY\",\n",
    "            \"value\": os.getenv(\"CLEARML_API_SECRET_KEY\")\n",
    "        },\n",
    "    ])\n",
    "\n",
    "if WANDB_LOGGED_IN:\n",
    "    if \"Envs\" not in drive_map.keys():\n",
    "        drive_map[\"Envs\"] = []\n",
    "    # Weights and biases currently requires access to the\n",
    "    # /.config directory in the docker. Therefore, the docker\n",
    "    # must be instantiated as root user. With the cells mentioned below\n",
    "    # we will be deleting the cells that set user roles.\n",
    "    if \"user\" in drive_map[\"DockerOptions\"].keys():\n",
    "        del(drive_map[\"DockerOptions\"][\"user\"])\n",
    "    drive_map[\"Envs\"].extend([\n",
    "        {\n",
    "            \"variable\": \"WANDB_API_KEY\",\n",
    "            \"value\": os.getenv(\"WANDB_API_KEY\")\n",
    "        }\n",
    "    ])\n",
    "\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(drive_map, mfile, indent=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Install the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher wheel.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will be using the kitti object detection dataset for this example. To find more details, please visit http://www.cvlibs.net/datasets/kitti/eval_object.php?obj_benchmark=2d. Please download both, the left color images of the object dataset from [here](http://www.cvlibs.net/download.php?file=data_object_image_2.zip) and, the training labels for the object dataset from [here](http://www.cvlibs.net/download.php?file=data_object_label_2.zip), and place the zip files in `$LOCAL_DATA_DIR`\n",
    "\n",
    "The data will then be extracted to have\n",
    "* training images in `$LOCAL_DATA_DIR/training/image_2`\n",
    "* training labels in `$LOCAL_DATA_DIR/training/label_2`\n",
    "* testing images in `$LOCAL_DATA_DIR/testing/image_2`\n",
    "\n",
    "You may use this notebook with your own dataset as well. To use this example with your own dataset, please follow the same directory structure as mentioned below.\n",
    "\n",
    "*Note: There are no labels for the testing images, therefore we use it just to visualize inferences for the trained model.*"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Download the dataset <a class=\"anchor\" id=\"head-2-1\"></a>\n",
    "Once you have gotten the download links in your email, please populate them in place of the `KITTI_IMAGES_DOWNLOAD_URL` and the `KITTI_LABELS_DOWNLOAD_URL`. This next cell, will download the data and place in `$LOCAL_DATA_DIR`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "!mkdir -p $LOCAL_DATA_DIR\n",
    "os.environ[\"URL_IMAGES\"]=KITTI_IMAGES_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then wget $URL_IMAGES -O $LOCAL_DATA_DIR/data_object_image_2.zip; else echo \"image archive already downloaded\"; fi \n",
    "os.environ[\"URL_LABELS\"]=KITTI_LABELS_DOWNLOAD_URL\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then wget $URL_LABELS -O $LOCAL_DATA_DIR/data_object_label_2.zip; else echo \"label archive already downloaded\"; fi "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Verify downloaded dataset <a class=\"anchor\" id=\"head-2-2\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_image_2.zip ]; then echo 'Image zip file not found, please download.'; else echo 'Found Image zip file.';fi\n",
    "!if [ ! -f $LOCAL_DATA_DIR/data_object_label_2.zip ]; then echo 'Label zip file not found, please download.'; else echo 'Found Labels zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# This may take a while: verify integrity of zip files \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_image_2.zip | cut -d ' ' -f 1 | grep -xq '^351c5a2aa0cd9238b50174a3a62b846bc5855da256b82a196431d60ff8d43617$' ; \\\n",
    "if test $? -eq 0; then echo \"images OK\"; else echo \"images corrupt, redownload!\" && rm -f $LOCAL_DATA_DIR/data_object_image_2.zip; fi \n",
    "!sha256sum $LOCAL_DATA_DIR/data_object_label_2.zip | cut -d ' ' -f 1 | grep -xq '^4efc76220d867e1c31bb980bbf8cbc02599f02a9cb4350effa98dbb04aaed880$' ; \\\n",
    "if test $? -eq 0; then echo \"labels OK\"; else echo \"labels corrupt, redownload!\" && rm -f $LOCAL_DATA_DIR/data_object_label_2.zip; fi "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack downloaded datasets to $DATA_DOWNLOAD_DIR.\n",
    "# The training images will be under $DATA_DOWNLOAD_DIR/training/image_2 and \n",
    "# labels will be under $DATA_DOWNLOAD_DIR/training/label_2.\n",
    "# The testing images will be under $DATA_DOWNLOAD_DIR/testing/image_2.\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_image_2.zip -d $LOCAL_DATA_DIR\n",
    "!unzip -u $LOCAL_DATA_DIR/data_object_label_2.zip -d $LOCAL_DATA_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "import os\n",
    "\n",
    "DATA_DIR = os.environ.get('LOCAL_DATA_DIR')\n",
    "num_training_images = len(os.listdir(os.path.join(DATA_DIR, \"training/image_2\")))\n",
    "num_training_labels = len(os.listdir(os.path.join(DATA_DIR, \"training/label_2\")))\n",
    "num_testing_images = len(os.listdir(os.path.join(DATA_DIR, \"testing/image_2\")))\n",
    "print(\"Number of images in the train/val set. {}\".format(num_training_images))\n",
    "print(\"Number of labels in the train/val set. {}\".format(num_training_labels))\n",
    "print(\"Number of images in the test set. {}\".format(num_testing_images))"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Prepare tf records from kitti format dataset <a class=\"anchor\" id=\"head-2-3\"></a>\n",
    "\n",
    "* Update the tfrecords spec file to take in your kitti format dataset\n",
    "* Create the tfrecords using the detectnet_v2 dataset_convert \n",
    "\n",
    "*Note: TfRecords only need to be generated once.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"TFrecords conversion spec file for kitti training\")\n",
    "!cat $LOCAL_SPECS_DIR/detectnet_v2_tfrecords_kitti_trainval.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Creating a new directory for the output tfrecords dump.\n",
    "print(\"Converting Tfrecords for kitti trainval dataset\")\n",
    "!mkdir -p $LOCAL_DATA_DIR/tfrecords && rm -rf $LOCAL_DATA_DIR/tfrecords/*\n",
    "!tao model detectnet_v2 dataset_convert \\\n",
    "                  -d $SPECS_DIR/detectnet_v2_tfrecords_kitti_trainval.txt \\\n",
    "                  -o $DATA_DOWNLOAD_DIR/tfrecords/kitti_trainval/kitti_trainval \\\n",
    "                  -r $USER_EXPERIMENT_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_DATA_DIR/tfrecords/kitti_trainval/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Download pre-trained model <a class=\"anchor\" id=\"head-2-4\"></a>\n",
    "Download the correct pretrained model from the NGC model registry for your experiment. Please note that for DetectNet_v2, the input is expected to be 0-1 normalized with input channels in RGB order. Therefore, for optimum results please download model templates from `nvidia/tao/pretrained_detectnet_v2`. The templates are now organized as version strings. For example, to download a resnet18 model suitable for detectnet please resolve to the ngc object shown as `nvidia/tao/pretrained_detectnet_v2:resnet18`. \n",
    "\n",
    "All other models are in BGR order expect input preprocessing with mean subtraction and input channels. Using them as pretrained weights may result in suboptimal performance.\n",
    "\n",
    "You may also use this notebook with the following purpose-built pretrained models \n",
    "* [PeopleNet](https://ngc.nvidia.com/catalog/models/nvidia:tao:peoplenet)\n",
    "* [TrafficCamNet](https://ngc.nvidia.com/catalog/models/nvidia:tao:trafficcamnet)\n",
    "* [DashCamNet](https://ngc.nvidia.com/catalog/models/nvidia:tao:dashcamnet)\n",
    "* [FaceDetect-IR](https://ngc.nvidia.com/catalog/models/nvidia:tao:facedetectir) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $LOCAL_PROJECT_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $LOCAL_PROJECT_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $LOCAL_PROJECT_DIR/ngccli\n",
    "!unzip -u \"$LOCAL_PROJECT_DIR/ngccli/$CLI\" -d $LOCAL_PROJECT_DIR/ngccli/\n",
    "!rm $LOCAL_PROJECT_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"LOCAL_PROJECT_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# List models available in the model registry.\n",
    "!ngc registry model list nvidia/tao/pretrained_detectnet_v2:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create the target destination to download the model.\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download the pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/pretrained_detectnet_v2:resnet18 \\\n",
    "    --dest $LOCAL_EXPERIMENT_DIR/pretrained_resnet18"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/pretrained_resnet18/pretrained_detectnet_v2_vresnet18"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Tfrecords for the train datasets\n",
    "    * To use the newly generated tfrecords, update the dataset_config parameter in the spec file at `$SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt` \n",
    "    * Update the fold number to use for evaluation. In case of random data split, please use fold `0` only\n",
    "    * For sequence-wise split, you may use any fold generated from the dataset convert tool\n",
    "* Pre-trained models\n",
    "* Augmentation parameters for on the fly data augmentation\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $LOCAL_SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "\n",
    "*Note: The training may take hours to complete. Also, the remaining notebook, assumes that the training was done in single-GPU mode. When run in multi-GPU mode, please expect to update the pruning and inference steps with new pruning thresholds and updated parameters in the clusterfile.json accordingly for optimum performance.*\n",
    "\n",
    "*Detectnet_v2 now supports restart from checkpoint. In case the training job is killed prematurely, you may resume training from the closest checkpoint by simply re-running the **same** command line. Please do make sure to use the <u>**same number of GPUs**</u> when restarting the training.*\n",
    "\n",
    "*When running the training with NUM_GPUs>1, you may need to modify the `batch_size_per_gpu` and `learning_rate` to get similar mAP as a 1GPU training run. In most cases, scaling down the batch-size by a factor of NUM_GPU's or scaling up the learning rate by a factor of NUM_GPU's would be a good place to start.* "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 train -e $SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt \\\n",
    "                        -r $USER_EXPERIMENT_DIR/experiment_dir_unpruned \\\n",
    "                        -n resnet18_detector \\\n",
    "                        --gpus $NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Model for each epoch:')\n",
    "print('---------------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/experiment_dir_unpruned/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate the trained model <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 evaluate -e $SPECS_DIR/detectnet_v2_train_resnet18_kitti.txt\\\n",
    "                           -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/resnet18_detector.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune the trained model <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Equalization criterion (`Applicable for resnets and mobilenets`)\n",
    "* Threshold for pruning.\n",
    "* Output directory to store the model\n",
    "\n",
    "*Usually, you just need to adjust `-pth` (threshold) for accuracy and model size trade off. Higher `pth` gives you smaller model (and thus higher inference speed) but worse accuracy. The threshold to use is dependent on the dataset. A pth value `5.2e-6` is just a start point. If the retrain accuracy is good, you can increase this value to get smaller models. Otherwise, lower this value to get better accuracy.*\n",
    "\n",
    "*For some internal studies, we have noticed that a pth value of 0.01 is a good starting point for detectnet_v2 models.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create an output directory if it doesn't exist.\n",
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 prune \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_unpruned/weights/resnet18_detector.hdf5 \\\n",
    "                  -o $USER_EXPERIMENT_DIR/experiment_dir_pruned/resnet18_nopool_bn_detectnet_v2_pruned.hdf5 \\\n",
    "                  -eq union \\\n",
    "                  -pth 0.0000052"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_pruned/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain the pruned model <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification with pretrained weights as pruned model.\n",
    "\n",
    "*Note: For retraining, please set the `load_graph` option to `true` in the model_config to load the pruned model graph. Also, if after retraining, the model shows some decrease in mAP, it could be that the originally trained model was pruned a little too much. Please try reducing the pruning threshold (thereby reducing the pruning ratio) and use the new model to retrain.*\n",
    "\n",
    "*Note: DetectNet_v2 now supports Quantization Aware Training, to help with optmizing the model. By default, the training in the cell below doesn't run the model with QAT enabled. For information on training a model with QAT, please refer to the cells under [section 11](#head-11)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Printing the retrain experiment file. \n",
    "# Note: We have updated the experiment file to include the \n",
    "# newly pruned model as a pretrained weights and, the\n",
    "# load_graph option is set to true \n",
    "!cat $LOCAL_SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model detectnet_v2 train -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt \\\n",
    "                        -r $USER_EXPERIMENT_DIR/experiment_dir_retrain \\\n",
    "                        -n resnet18_detector_pruned \\\n",
    "                        --gpus $NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate the retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This section evaluates the pruned and retrained model, using the `evaluate` command."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 evaluate -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt \\\n",
    "                           -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/resnet18_detector_pruned.hdf5"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Visualize inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    "In this section, we run the `inference` tool to generate inferences on the trained models. To render bboxes from more classes, please edit the spec file `detectnet_v2_inference_kitti_tlt.txt` to include all the classes you would like to visualize and edit the rest of the file accordingly."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy some test images\n",
    "!mkdir -p $LOCAL_DATA_DIR/test_samples\n",
    "!cp $LOCAL_DATA_DIR/testing/image_2/00000* $LOCAL_DATA_DIR/test_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "!tao model detectnet_v2 inference -e $SPECS_DIR/detectnet_v2_inference_kitti_tlt.txt \\\n",
    "                            -r $USER_EXPERIMENT_DIR/tlt_infer_testing \\\n",
    "                            -i $DATA_DOWNLOAD_DIR/test_samples"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The `inference` tool produces two outputs. \n",
    "1. Overlain images in `$USER_EXPERIMENT_DIR/tlt_infer_testing/images_annotated`\n",
    "2. Frame by frame bbox labels in kitti format located in `$USER_EXPERIMENT_DIR/tlt_infer_testing/labels`\n",
    "\n",
    "*Note: To run inferences for a single image, simply replace the path to the -i flag in `inference` command with the path to the image.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Simple grid visualizer\n",
    "!pip3 install \"matplotlib>=3.3.3, <4.0\"\n",
    "%matplotlib inline\n",
    "import matplotlib.pyplot as plt\n",
    "import os\n",
    "from math import ceil\n",
    "valid_image_ext = ['.jpg', '.png', '.jpeg', '.ppm']\n",
    "\n",
    "def visualize_images(image_dir, num_cols=4, num_images=10):\n",
    "    output_path = os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'], image_dir)\n",
    "    num_rows = int(ceil(float(num_images) / float(num_cols)))\n",
    "    f, axarr = plt.subplots(num_rows, num_cols, figsize=[80,30])\n",
    "    f.tight_layout()\n",
    "    a = [os.path.join(output_path, image) for image in os.listdir(output_path) \n",
    "         if os.path.splitext(image)[1].lower() in valid_image_ext]\n",
    "    for idx, img_path in enumerate(a[:num_images]):\n",
    "        col_id = idx % num_cols\n",
    "        row_id = idx // num_cols\n",
    "        img = plt.imread(img_path)\n",
    "        axarr[row_id, col_id].imshow(img) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Visualizing the first 12 images.\n",
    "OUTPUT_PATH = 'tlt_infer_testing/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 4 # number of columns in the visualizer grid.\n",
    "IMAGES = 12 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Model Export <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $LOCAL_EXPERIMENT_DIR/experiment_dir_final\n",
    "# Removing a pre-existing copy of the onnx if there has been any.\n",
    "import os\n",
    "output_file=os.path.join(os.environ['LOCAL_EXPERIMENT_DIR'],\n",
    "                         \"experiment_dir_final/resnet18_detector.onnx\")\n",
    "if os.path.exists(output_file):\n",
    "    os.system(\"rm {}\".format(output_file))\n",
    "!tao model detectnet_v2 export \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_retrain/weights/resnet18_detector_pruned.hdf5 \\\n",
    "                  -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt \\\n",
    "                  -o $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.onnx \\\n",
    "                  --onnx_route tf2onnx \\\n",
    "                  --gen_ds_config"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference. \n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's. \n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson from the dev zone link [here](https://developer.nvidia.com/tao-converter). \n",
    "\n",
    "If you choose to integrate your model into deepstream directly, you may do so by simply copying the exported `.onnx` file along with the calibration cache to the target device and updating the spec file that configures the `gst-nvinfer` element to point to this newly exported model. Usually this file is called `config_infer_primary.txt` for detection models and `config_infer_secondary_*.txt` for classification models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Need to pass the actual image directory instead of data root for tao deploy to locate images for calibration\n",
    "!sed -i \"s|/workspace/tao-experiments/data/training|/workspace/tao-experiments/data/training/image_2|g\" $LOCAL_SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt\n",
    "# Convert to TensorRT engine (INT8)\n",
    "!tao deploy detectnet_v2 gen_trt_engine \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.onnx \\\n",
    "                  --data_type int8 \\\n",
    "                  --batches 10 \\\n",
    "                  --batch_size 4 \\\n",
    "                  --max_batch_size 64\\\n",
    "                  --engine_file $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.trt.int8 \\\n",
    "                  --cal_cache_file $USER_EXPERIMENT_DIR/experiment_dir_final/calibration.bin \\\n",
    "                  -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt \\\n",
    "                  --results_dir $USER_EXPERIMENT_DIR/experiment_dir_final \\\n",
    "                  --verbose\n",
    "# Convert back the spec file\n",
    "!sed -i \"s|/workspace/tao-experiments/data/training/image_2|/workspace/tao-experiments/data/training|g\" $LOCAL_SPECS_DIR/detectnet_v2_retrain_resnet18_kitti.txt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "*Note: For this example, we generate a calibration tensorfile containing 10 batches of training data.\n",
    "Ideally, it is best to use at least 10-20% of the training data to do so. The more data provided during calibration, the closer int8 inferences are to fp32 inferences.*\n",
    "\n",
    "*Note: If the model was trained with QAT nodes available, please refrain from using the post training int8 optimization as mentioned below. Please export the model in int8 mode (using the arg `--data_type int8`) with just the path to the calibration cache file (using the argument `--cal_cache_file`) and calibration json file (using the argument `--cal_json_file`)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $LOCAL_EXPERIMENT_DIR/experiment_dir_final"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Verify Deployed Model <a class=\"anchor\" id=\"head-11\"></a>\n",
    "Verify the exported model by visualizing inferences on TensorRT.\n",
    "In addition to running inference on a `.hdf5` model in [step 9](#head-9), the `inference` tool is also capable of consuming the converted `TensorRT engine` from [step 10.B](#head-10-2).\n",
    "\n",
    "*If after int-8 calibration the accuracy of the int-8 inferences seem to degrade, it could be because the there wasn't enough data in the calibration tensorfile used to calibrate thee model or, the training data is not entirely representative of your test images, and the calibration maybe incorrect. Therefore, you may either regenerate the calibration tensorfile with more batches of the training data and recalibrate the model, or calibrate the model on a few images from the test set. This may be done using `--cal_image_dir` flag in the `tao deploy <task> gen_trt_engine` tool. For more information, please follow the instructions in the USER GUIDE."
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Inference using TensorRT engine <a class=\"anchor\" id=\"head-11-1\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy detectnet_v2 inference -e $SPECS_DIR/detectnet_v2_inference_kitti_etlt.txt \\\n",
    "                                   -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector.trt.int8 \\\n",
    "                                   -r $USER_EXPERIMENT_DIR/etlt_infer_testing \\\n",
    "                                   -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                                   -b 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the first 12 inferenced images.\n",
    "OUTPUT_PATH = 'etlt_infer_testing/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 4 # number of columns in the visualizer grid.\n",
    "IMAGES = 12 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. QAT workflow <a class=\"anchor\" id=\"head-12\"></a>\n",
    "This section delves into the newly enabled Quantization Aware Training feature with DetectNet_v2. The workflow defined below converts a pruned model from section [5](#head-5) to enable QAT and retrain this model to while accounting the noise introduced due to quantization in the forward pass. "
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### A. Convert pruned model to QAT and retrain <a class=\"anchor\" id=\"head-12-1\"></a>\n",
    "All detectnet models, unpruned and pruned models can be converted to QAT models by setting the `enable_qat` parameter in the `training_config` component of the spec file to `true`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Printing the retrain experiment file. \n",
    "# Note: We have updated the experiment file to convert the\n",
    "# pretrained model to qat mode by setting the enable_qat\n",
    "# parameter.\n",
    "!cat $LOCAL_SPECS_DIR/detectnet_v2_retrain_resnet18_kitti_qat.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 train -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti_qat.txt \\\n",
    "                        -r $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat \\\n",
    "                        -n resnet18_detector_pruned_qat \\\n",
    "                        --gpus $NUM_GPUS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $LOCAL_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### B. Evaluate QAT converted model <a class=\"anchor\" id=\"head-12-2\"></a>\n",
    "This section evaluates a QAT enabled pruned retrained model. The mAP of this model should be comparable to that of the pruned retrained model without QAT. However, due to quantization, it is possible sometimes to see a drop in the mAP value for certain datasets."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model detectnet_v2 evaluate -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti_qat.txt \\\n",
    "                           -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/resnet18_detector_pruned_qat.hdf5 \\\n",
    "                           -f tlt"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### C. Export QAT trained model to int8 <a class=\"anchor\" id=\"head-12-3\"></a>\n",
    "Export a QAT trained model to TensorRT parsable model. This command generates an .onnx file from the trained model and the serializes corresponding int8 scales as a TRT readable calibration cache file."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!rm -rf $LOCAL_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_qat.onnx\n",
    "!rm -rf $LOCAL_EXPERIMENT_DIR/experiment_dir_final/calibration_qat.bin\n",
    "!tao model detectnet_v2 export \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/weights/resnet18_detector_pruned_qat.hdf5 \\\n",
    "                  -o $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_qat.onnx \\\n",
    "                  -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti_qat.txt \\\n",
    "                  --cal_json_file $USER_EXPERIMENT_DIR/experiment_dir_final/calibration_qat.json \\\n",
    "                  --gen_ds_config \\\n",
    "                  --onnx_route tf2onnx \\\n",
    "                  --verbose"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy detectnet_v2 gen_trt_engine \\\n",
    "                  -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_qat.onnx \\\n",
    "                  -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti_qat.txt \\\n",
    "                  --data_type int8 \\\n",
    "                  --batch_size 64 \\\n",
    "                  --max_batch_size 64\\\n",
    "                  --engine_file $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_qat.trt.int8 \\\n",
    "                  --cal_cache_file $USER_EXPERIMENT_DIR/experiment_dir_final/calibration_qat.bin \\\n",
    "                  --cal_json_file $USER_EXPERIMENT_DIR/experiment_dir_final/calibration_qat.json \\\n",
    "                  --results_dir $USER_EXPERIMENT_DIR/experiment_dir_final \\\n",
    "                  --verbose"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### D. Evaluate a QAT trained model using the exported TensorRT engine <a class=\"anchor\" id=\"head-12-4\"></a>\n",
    "This section evaluates a QAT enabled pruned retrained model using the TensorRT int8 engine that was exported in [Section C](#head-12-3). Please note that there maybe a slight difference (~0.1-0.5%) in the mAP from [Section B](#head-12-2), oweing to some differences in the implementation of quantization in TensorRT.\n",
    "\n",
    "*Note: The TensorRT evaluator might be slightly slower than the TAO evaluator here, because the evaluation dataloader is pinned to the CPU to avoid any clashes between TensorRT and TAO instances in the GPU. Please note that this tool was not intended and has not been developed for profiling the model. It is just a means to qualitatively analyse the model.*\n",
    "\n",
    "*Please use native TensorRT or DeepStream for the most optimized inferences.*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import shutil\n",
    "val_split = 0.14  # Val split from the spec file\n",
    "img_ext = '.png'\n",
    "lab_ext = '.txt'\n",
    "\n",
    "# Get images and labels from the partitioned validation set\n",
    "images_root = f\"{os.environ['LOCAL_DATA_DIR']}/training/image_2\"\n",
    "labels_root = f\"{os.environ['LOCAL_DATA_DIR']}/training/label_2\"\n",
    "images_list = [os.path.splitext(imfile)[0] for imfile in sorted(os.listdir(images_root)) if imfile.endswith(img_ext)]\n",
    "num_val_images = int(len(images_list) * val_split)\n",
    "\n",
    "# Copy the data to a separate directory for evaluation\n",
    "os.makedirs(os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"images_val\"), exist_ok=True)\n",
    "os.makedirs(os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"labels_val\"), exist_ok=True)\n",
    "\n",
    "for fname in images_list[:num_val_images]:\n",
    "    shutil.copyfile(os.path.join(images_root, fname + img_ext), \\\n",
    "                             os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"images_val\", fname + img_ext))\n",
    "    shutil.copyfile(os.path.join(labels_root, fname + lab_ext), \\\n",
    "                             os.path.join(f\"{os.environ['LOCAL_DATA_DIR']}/training\", \"labels_val\", fname + lab_ext))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy detectnet_v2 evaluate -e $SPECS_DIR/detectnet_v2_retrain_resnet18_kitti_qat.txt \\\n",
    "                                  -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_qat.trt.int8 \\\n",
    "                                  -i $DATA_DOWNLOAD_DIR/training/images_val \\\n",
    "                                  -l $DATA_DOWNLOAD_DIR/training/labels_val \\\n",
    "                                  -r $USER_EXPERIMENT_DIR/experiment_dir_retrain_qat/"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### E. Inference using QAT engine <a class=\"anchor\" id=\"head-12-5\"></a>\n",
    "Run inference and visualize detections on test images, using the exported TensorRT engine from [Section C](#head-12-3)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao deploy detectnet_v2 inference -e $SPECS_DIR/detectnet_v2_inference_kitti_etlt_qat.txt \\\n",
    "                                   -m $USER_EXPERIMENT_DIR/experiment_dir_final/resnet18_detector_qat.trt.int8 \\\n",
    "                                   -r $USER_EXPERIMENT_DIR/tlt_infer_testing_qat \\\n",
    "                                   -i $DATA_DOWNLOAD_DIR/test_samples \\\n",
    "                                   -b 4"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# visualize the first 12 inferenced images.\n",
    "OUTPUT_PATH = 'tlt_infer_testing_qat/images_annotated' # relative path from $USER_EXPERIMENT_DIR.\n",
    "COLS = 4 # number of columns in the visualizer grid.\n",
    "IMAGES = 12 # number of images to visualize.\n",
    "\n",
    "visualize_images(OUTPUT_PATH, num_cols=COLS, num_images=IMAGES)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "f23a2831654361cfd8b219e05b5055fdda3e37fe5c0b020e6226f740844c300a"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
